{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent Management Through AgentSet\n",
    "\n",
    "### The Boltzmann Wealth Model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to get straight to the tutorial checkout these environment providers:<br>\n",
    "(with Google Account) [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa/blob/main/docs/tutorials/3_agentset.ipynb)<br>\n",
    "(No Google Account) [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/projectmesa/mesa/main?labpath=docs%2Ftutorials%2F3_agentset.ipynb) (This can take 30 seconds to 5 minutes to load)\n",
    "\n",
    "*If you are running locally, please ensure you have the latest Mesa version installed.*\n",
    "\n",
    "## Tutorial Description\n",
    "\n",
    "This tutorial extends the Boltzmann wealth model from the [Collecting Data tutorial](https://mesa.readthedocs.io/latest/tutorials/2_collecting_data.html), by demonstrating Mesa's AgentSet functionality. \n",
    "\n",
    "In this portion, we will demonstrate how users can employ AgentSet for different purposes. \n",
    "\n",
    "*If you are starting here please see the [Running Your First Model tutorial](https://mesa.readthedocs.io/latest/tutorials/0_first_model.html) for dependency and start-up instructions*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IN COLAB? - Run the next cell "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "%pip install --quiet mesa[rec]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Dependencies\n",
    "This includes importing of dependencies needed for the tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Has multi-dimensional arrays and matrices.\n",
    "# Has a large collection of mathematical functions to operate on these arrays.\n",
    "import numpy as np\n",
    "\n",
    "# Data manipulation and analysis.\n",
    "import pandas as pd\n",
    "\n",
    "# Data visualization tools.\n",
    "import seaborn as sns\n",
    "\n",
    "import mesa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent Management Through AgentSet\n",
    "\n",
    "**Background:** Mesa uses a set based approach, [AgentSet](https://github.com/projectmesa/mesa/blob/f511a4bc57340cb2dd0ba4b0af76307b37aea0ca/mesa/agent.py#L147) to allow users efficiently and intuitively manage their agents.  For the most part users will never explicitly call AgentSet and in fact, we have already used the AgentSet methods functionality when we used `shuffle_do(move)` to reorder the agents and then `do(exchange)` to have the agents exchange money in sequence. Although you will likely never interact with AgentSent directly it is important to know the Mesa uses a set based approach for agent management. \n",
    "\n",
    "Beyond the method functionality there are additional ways AgentSet can help you manage your agents and we will look at two additional examples in this tutorial, but you can see more in the [Getting Started Section of Mesa](https://mesa.readthedocs.io/stable/getting_started.html#agentset-functionality). \n",
    "\n",
    "**Model-specific information:** We will show two agent management techniques just to demonstrate the capability\n",
    "1. **Selecting** We will institute a policy that has the rich agents give money to the  poor agents\n",
    "2. **GroupBy** We will group agents together based on wealth\n",
    "\n",
    "*A big thanks to @Ewout for his exceptional work on developing and implementing AgentSet*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selecting \n",
    "\n",
    "**Model-specific Information:** For this variation of the model we are going to institute a policy that only rich agents give money to poor agent\n",
    "\n",
    "**Code Implementation:** We will use `agents.select` to separate the agents into rich and poor agents. If there are rich agents then they are the only ones who give money. \n",
    "\n",
    "\\# Get lists of rich and poor agents\n",
    "\n",
    "- **Description:** Uses `AgentSet.select` with a function (in this case a lambda function) to select agents with greater than 3 units of wealth and less than three units of wealth. This will give us two lists of agents rich agent and poor agent which we can then use to execute the `give_money` method. \n",
    "- **API:** [AgentSet.select](https://mesa.readthedocs.io/latest/apis/agent.html#mesa.agent.AgentSet.select)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_gini(model):\n",
    "    agent_wealths = [agent.wealth for agent in model.agents]\n",
    "    x = sorted(agent_wealths)\n",
    "    n = model.num_agents\n",
    "    B = sum(xi * (n - i) for i, xi in enumerate(x)) / (n * sum(x))\n",
    "    return 1 + (1 / n) - 2 * B\n",
    "\n",
    "\n",
    "class MoneyAgent(mesa.Agent):\n",
    "    \"\"\"An agent with fixed initial wealth.\"\"\"\n",
    "\n",
    "    def __init__(self, model):\n",
    "        super().__init__(model)\n",
    "        self.wealth = 1\n",
    "\n",
    "    def give_money(self, poor_agents):\n",
    "        if self.wealth > 0:\n",
    "            other_agent = self.random.choice(poor_agents)\n",
    "            other_agent.wealth += 1\n",
    "            self.wealth -= 1\n",
    "\n",
    "\n",
    "class MoneyModel(mesa.Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "\n",
    "    def __init__(self, n):\n",
    "        super().__init__()\n",
    "        self.num_agents = n\n",
    "\n",
    "        # Create agents\n",
    "        MoneyAgent.create_agents(model=self, n=n)\n",
    "\n",
    "        self.datacollector = mesa.DataCollector(\n",
    "            model_reporters={\"Gini\": compute_gini}, agent_reporters={\"Wealth\": \"wealth\"}\n",
    "        )\n",
    "\n",
    "    def step(self):\n",
    "        self.datacollector.collect(self)\n",
    "        # Get lists of rich and poor agents\n",
    "        rich_agents = model.agents.select(lambda a: a.wealth >= 3)\n",
    "        poor_agents = model.agents.select(lambda a: a.wealth < 3)\n",
    "        # When there is rich agents only have them give money to poor agents\n",
    "        if len(rich_agents) > 0:\n",
    "            rich_agents.shuffle_do(\"give_money\", poor_agents)\n",
    "        else:\n",
    "            poor_agents.shuffle_do(\"give_money\", poor_agents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now run the model, collect the data, and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MoneyModel(100)\n",
    "for _ in range(20):\n",
    "    model.step()\n",
    "\n",
    "\n",
    "data = model.datacollector.get_agent_vars_dataframe()\n",
    "# Use seaborn\n",
    "g = sns.histplot(data[\"Wealth\"], discrete=True)\n",
    "g.set(title=\"Wealth distribution\", xlabel=\"Wealth\", ylabel=\"number of agents\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Group By \n",
    "\n",
    "**Model-specific implementation:** In this case we will give agents an attribute of ethnicity of Green, Blue or Mixed. Green and Blue agents only give money to their ethnicity while Mixed can give money to anyone.\n",
    "\n",
    "**Code Implementation**: Using `groupby` we will execute the above logic in our code by passing a list of grouped agents into our `give_money` function.  To ensure we can plot wealth by group we also need to add ethnicity to our datacollector. \n",
    "\n",
    "\\# Create dictionary of agents groupby\n",
    "\n",
    "**Description:** Uses `AgentSet.groupby` to group agents by their ethnicity attribute. This will give us a dictionary where the keys are the different ethnicities and the values are an `AgentSet`. In this case we will then use the `AgentSet` class and leverage its `shuffle_do` capability to then give money to the target groups. \n",
    "- **API:** [AgentSet.select](https://mesa.readthedocs.io/latest/apis/agent.html#mesa.agent.AgentSet.groupby)\n",
    "- **Note:** `AgentSet` has a lot of functionality and similar to `discrete_space` has the ability to add new features and make Mesa models more user-friendly. We strongly encourage you to check out the [AgentSet API](https://mesa.readthedocs.io/latest/apis/agent.html#mesa.agent.AgentSet) to see all the functionality and if you have an idea feel free to [contribute](https://github.com/projectmesa/mesa/blob/main/CONTRIBUTING.md) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoneyAgent(mesa.Agent):\n",
    "    \"\"\"An agent with fixed initial wealth.\"\"\"\n",
    "\n",
    "    def __init__(self, model, ethnicity):\n",
    "        super().__init__(model)\n",
    "        self.wealth = 1\n",
    "        self.ethnicity = ethnicity\n",
    "\n",
    "    def give_money(self, similars):\n",
    "        if self.wealth > 0:\n",
    "            other_agent = self.random.choice(similars)\n",
    "            other_agent.wealth += 1\n",
    "            self.wealth -= 1\n",
    "\n",
    "\n",
    "class MoneyModel(mesa.Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "\n",
    "    def __init__(self, n):\n",
    "        super().__init__()\n",
    "        self.num_agents = n\n",
    "\n",
    "        # Create a list of our different ethnicities\n",
    "        ethnicities = [\"Green\", \"Blue\", \"Mixed\"]\n",
    "\n",
    "        # Create agents\n",
    "        MoneyAgent.create_agents(\n",
    "            model=self,\n",
    "            n=self.num_agents,\n",
    "            ethnicity=self.random.choices(ethnicities, k=self.num_agents),\n",
    "        )\n",
    "\n",
    "        self.datacollector = mesa.DataCollector(\n",
    "            model_reporters={\"Gini\": compute_gini},\n",
    "            agent_reporters={\"Wealth\": \"wealth\", \"Ethnicity\": \"ethnicity\"},\n",
    "        )\n",
    "\n",
    "    def step(self):\n",
    "        self.datacollector.collect(self)\n",
    "        # Create dictionary of agents groupby\n",
    "        grouped_agents = model.agents.groupby(\"ethnicity\")\n",
    "        for ethnic, similars in grouped_agents:\n",
    "            if ethnic != \"Mixed\":\n",
    "                similars.shuffle_do(\"give_money\", similars)\n",
    "            else:\n",
    "                similars.shuffle_do(\n",
    "                    \"give_money\", self.agents\n",
    "                )  # This allows mixed to trade with anyone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the model\n",
    "model = MoneyModel(100)\n",
    "for _ in range(20):\n",
    "    model.step()\n",
    "\n",
    "# get the data\n",
    "data = model.datacollector.get_agent_vars_dataframe()\n",
    "# assign histogram colors\n",
    "palette = {\"Green\": \"green\", \"Blue\": \"blue\", \"Mixed\": \"purple\"}\n",
    "g = sns.histplot(data=data, x=\"Wealth\", hue=\"Ethnicity\", discrete=True, palette=palette)\n",
    "g.set(title=\"Wealth distribution\", xlabel=\"Wealth\", ylabel=\"number of agents\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises \n",
    "- Create a new policy or alter an existing policy in this model to see the impact\n",
    "- Use a different feature in `AgentSet` and integrate into this model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "Check out the [basic visualization tutorial](https://mesa.readthedocs.io/latest/tutorials/4_visualization_basic.html) on how to build interactive dashboards for your models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Comer2014] Comer, Kenneth W. “Who Goes First? An Examination of the Impact of Activation on Outcome Behavior in AgentBased Models.” George Mason University, 2014. http://mars.gmu.edu/bitstream/handle/1920/9070/Comer_gmu_0883E_10539.pdf\n",
    "\n",
    "[Dragulescu2002] Drăgulescu, Adrian A., and Victor M. Yakovenko. “Statistical Mechanics of Money, Income, and Wealth: A Short Survey.” arXiv Preprint Cond-mat/0211175, 2002. http://arxiv.org/abs/cond-mat/0211175."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
